import base64
import hashlib
import os
import time
import requests
from rembg import remove
from typing import Union, List, Dict, Tuple, Optional
from PIL import Image
from io import BytesIO
from langchain_openai.chat_models.base import BaseChatModel
from loguru import logger

from ..config.annotation import print_time

@print_time
def generate_image_file(llm: BaseChatModel,
                        image_desc: str,
                        image_size: Union[str, List[str]] = "1024x1024",
                        input_images: Union[str, List[str]] = None,
                        valid_image_sizes: List[str] = ["1024x1024","768x1344","864x1152","1344x768","1152x864","1440x720","720x1440"],
                        image_type: str = "png",
                        return_full_path: bool = False) -> Union[Dict[str, str], Tuple[Dict[str, str], Dict], None]:
    """
    生成图片文件并保存到磁盘

    Args:
        llm: 用于图片生成的语言模型
        image_desc: 图片描述
        image_size: 图片尺寸
        input_images: 用于图片到图片生成的输入图片 (URL 或 base64)
        valid_image_sizes: 模型支持的有效尺寸列表
        image_type: 图片文件类型 (png, jpg, 等)
        return_full_path: 是否返回完整文件路径

    Returns:
        (image_files_dict, usage_metadata) 的元组或仅 image_files_dict 以保持向后兼容性
    """
    try:
        logger.debug(
            f"Generating image file - desc: '{image_desc[:30]}...', size: {image_size}, type: {image_type}")

        result = generate_image_base64(llm, image_desc, image_size, input_images, valid_image_sizes)
        if not result:
            logger.warning(f"Failed to generate base64 image for: '{image_desc[:30]}...'")
            return None

        # Handle both old and new return formats
        if isinstance(result, tuple):
            image_base64_result, usage_metadata = result
        else:
            image_base64_result, usage_metadata = result, {}

        def save_single_image(base64_str: str, size: str = "") -> str:
            timestamp = int(time.time() * 1000)
            filename = f"{timestamp}_{size}.{image_type}".strip("_")
            filepath = os.path.join(get_or_create_dir(filename=filename), filename)
            img_data = base64.b64decode(base64_str)
            with open(filepath, 'wb') as f:
                f.write(img_data)
            logger.info(f"Image saved{' (size: ' + size + ')' if size else ''}: {filepath}")
            return filepath if return_full_path else filename

        if isinstance(image_base64_result, str):
            size = image_size[0] if isinstance(image_size, list) else image_size
            files_result = {size: save_single_image(image_base64_result, size)}
        else:
            files_result = {size: save_single_image(base64_str, size)
                            for size, base64_str in image_base64_result.items()}

        if usage_metadata:
            return files_result, usage_metadata
        return files_result

    except Exception as e:
        logger.error(f"Image generation failed: {str(e)}")
        return None


def generate_image_base64(llm: BaseChatModel,
                          image_desc: str,
                          image_size: Union[str, List[str]] = "1024x1024",
                          input_images: Union[str, List[str]] = None,
                          valid_image_sizes: List[str] = ["1024x1024","768x1344","864x1152","1344x768","1152x864","1440x720","720x1440"]) -> Union[
    Tuple[Union[str, Dict[str, str]], Dict], None]:
    """
    生成base64编码的图片

    Args:
        llm: 用于图片生成的语言模型
        image_desc: 图片描述
        image_size: 图片尺寸
        input_images: 用于图片到图片生成的输入图片 (URL 或 base64)
        valid_image_sizes: 模型支持的有效尺寸列表

    Returns:
        (图片, usage_metadata) 的元组
    """
    if not llm:
        logger.error("No LLM provided for image generation")
        return None

    try:
        # Check if this is a Gemini model that needs chat interface
        extra_body = getattr(llm, 'extra_body', {})
        if isinstance(extra_body, dict) and extra_body.get('type') == 'gemini':
            return generate_image_via_chat(llm, image_desc, image_size, input_images)

        image_sizes = [image_size] if isinstance(image_size, str) else image_size

        def parse_size(size_str: str) -> tuple:
            """解析尺寸字符串，处理 'WxH' 和 'W*H' 格式"""
            if 'x' in size_str:
                return tuple(map(int, size_str.split('x')))
            elif '*' in size_str:
                return tuple(map(int, size_str.split('*')))
            else:
                raise ValueError(f"Invalid size format: {size_str}")

        target_size = max(image_sizes, key=lambda x: parse_size(x)[0] * parse_size(x)[1])

        if valid_image_sizes:
            # Parse target dimensions
            target_w, target_h = parse_size(target_size)
            target_ratio = target_w / target_h

            # First check for exact match
            if target_size in valid_image_sizes:
                target_size = target_size
            else:
                # Look for same aspect ratio
                same_ratio_sizes = []
                for size in valid_image_sizes:
                    w, h = parse_size(size)
                    ratio = w / h
                    if abs(ratio - target_ratio) < 0.001:  # Account for float precision
                        same_ratio_sizes.append(size)

                if same_ratio_sizes:
                    # Pick the closest resolution
                    target_size = min(same_ratio_sizes,
                                      key=lambda x: abs(parse_size(x)[0] - target_w))
                else:
                    # Pick size with closest ratio
                    target_size = min(valid_image_sizes,
                                      key=lambda x: abs(parse_size(x)[0] / parse_size(x)[1] - target_ratio))

        # Prepare extra_body with input images if provided
        request_extra_body = dict(llm.extra_body) if llm.extra_body else {}

        if input_images:
            images = input_images if isinstance(input_images, list) else [input_images]

            processed_images = []
            for img in images:
                if not img:
                    continue
                if img.startswith(('http://', 'https://', 'data:image/')):
                    processed_images.append(img)
                else:
                    image_type = detect_image_type(img)
                    processed_images.append(f"data:image/{image_type};base64,{img}")

            if processed_images:
                request_extra_body["image"] = processed_images if len(processed_images) > 1 else processed_images[0]

        result = llm.root_client.images.generate(
            model=llm.model_name,
            prompt=image_desc,
            size=target_size,
            response_format="b64_json",
            extra_body=request_extra_body
        )

        usage_metadata = {}
        if hasattr(result, 'usage') and result.usage:
            usage_metadata = dict(result.usage)
            usage_metadata['model'] = llm.model_name

        base_output = result.data[0].b64_json if hasattr(result, 'data') else None
        if not base_output:
            logger.warning("Image generation returned no base64 data")
            return None

        # If only one size requested, return string
        if len(image_sizes) == 1:
            if target_size != image_sizes[0]:
                logger.debug(f"Resizing from {target_size} to {image_sizes[0]}")
                resized_output = resize_base64_image(base_output, image_sizes[0])
                return resized_output, usage_metadata
            return base_output, usage_metadata

        # Generate all requested sizes
        outputs = {}
        for size in image_sizes:
            if size == target_size:
                outputs[size] = base_output
            else:
                logger.debug(f"Resizing from {target_size} to {size}")
                outputs[size] = resize_base64_image(base_output, size)
        return outputs, usage_metadata

    except Exception as e:
        logger.error(f"Image generation error: {str(e)}")
        return None


def generate_image_via_chat(llm: BaseChatModel,
                            image_desc: str,
                            image_size: Union[str, List[str]] = "1024x1024",
                            input_images: Union[str, List[str]] = None) -> Union[str, Dict[str, str], Tuple[Union[str, Dict[str, str]], Dict], None]:
    """
    使用 OpenAI 聊天完成 API 为 Gemini 模型生成图片

    Args:
        llm: 语言模型实例
        image_desc: 图片描述
        image_size: 目标图片尺寸
        input_images: 用于图片到图片生成的输入图片 (URL 或 base64)

    Returns:
        Base64 编码的图片字符串、尺寸到 base64 字符串的字典，或包含使用元数据的元组
    """
    try:
        messages = [{"role": "user", "content": [{"type": "text", "text": f"Generate an image: {image_desc}"}]}]

        if input_images:
            images = input_images if isinstance(input_images, list) else [input_images]

            for img in images:
                if not img:
                    continue

                if img.startswith(('http://', 'https://', 'data:image/')):
                    image_url = img
                else:
                    image_type = detect_image_type(img)
                    image_url = f"data:image/{image_type};base64,{img}"

                messages[0]["content"].append({
                    "type": "image_url",
                    "image_url": {
                        "url": image_url
                    }
                })

        response = llm.root_client.chat.completions.create(
            model=llm.model_name,
            messages=messages,
            stream=False,
            extra_body=llm.extra_body
        )
        usage_metadata = {}
        if hasattr(response, 'usage') and response.usage:
            usage_metadata = dict(response.usage)
            usage_metadata['model'] = llm.model_name
        usage_metadata['generated_images'] = 1

        if not response.choices or not hasattr(response.choices[0].message, 'images') or not response.choices[
            0].message.images:
            logger.warning("No images in chat response")
            return None

        image_url = response.choices[0].message.images[0]["image_url"]["url"]
        if not image_url.startswith('data:image/'):
            logger.warning("Invalid image URL format")
            return None

        base64_data = image_url.split(',', 1)[1] if ',' in image_url else image_url
        image_sizes = [image_size] if isinstance(image_size, str) else image_size

        if len(image_sizes) == 1:
            return (base64_data, usage_metadata) if usage_metadata else base64_data

        outputs = {size: resize_base64_image(base64_data, size) for size in image_sizes}
        return (outputs, usage_metadata) if usage_metadata else outputs

    except Exception as e:
        logger.error(f"Chat-based image generation error: {str(e)}")
        return None


def resize_base64_image(base64_str: str, target_size: str, keep_ratio: bool = False) -> str:
    """
    将 base64 编码的图片调整到目标尺寸

    Args:
        base64_str: Base64 编码的图片字符串
        target_size: 目标尺寸，格式为 "widthxheight" 或 "width*height"
        keep_ratio: 是否保持宽高比

    Returns:
        Base64 编码的调整后图片字符串
    """
    try:
        # Decode base64 to image
        img_data = base64.b64decode(base64_str)
        img = Image.open(BytesIO(img_data))

        # Parse target size, handling both 'WxH' and 'W*H' formats
        if 'x' in target_size:
            target_w, target_h = map(int, target_size.split('x'))
        elif '*' in target_size:
            target_w, target_h = map(int, target_size.split('*'))
        else:
            raise ValueError(f"Invalid size format: {target_size}")

        if keep_ratio:
            # Maintain aspect ratio using thumbnail method
            img.thumbnail((target_w, target_h), Image.Resampling.LANCZOS)
        else:
            # Force resize to exact dimensions
            img = img.resize((target_w, target_h), Image.Resampling.LANCZOS)

        # Convert back to base64
        buffered = BytesIO()
        img.save(buffered, format=img.format or "PNG")
        return base64.b64encode(buffered.getvalue()).decode()
    except Exception as e:
        logger.error(f"Image resize error: {str(e)}")
        return base64_str


def get_base64_file_size(base64_str) -> int:
    """
    获取 base64 编码文件的字节大小

    Args:
        base64_str: Base64 编码字符串

    Returns:
        字节大小
    """
    if ',' in base64_str:
        base64_str = base64_str.split(',')[1]

    padding = len(base64_str) % 4
    if padding:
        base64_str += '=' * (4 - padding)
    return len(base64.b64decode(base64_str))


def get_or_create_dir(filename: str) -> str:
    """
    基于文件名哈希获取或创建目录路径

    Args:
        filename: 要哈希的文件名

    Returns:
        目录路径
    """
    h = md5_encode(filename)
    root_path = os.getenv("TMP_PATH", "/tmp")
    path = f"{root_path}/{h[1:3]}/{h[4:6]}"
    if not os.path.exists(path):
        os.makedirs(path)
    return path


def md5_encode(string) -> str:
    """
    创建字符串的 MD5 哈希

    Args:
        string: 要哈希的字符串

    Returns:
        十六进制字符串形式的 MD5 哈希
    """
    md5 = hashlib.md5()
    md5.update(string.encode('utf-8'))
    hex_digest = md5.hexdigest()
    return hex_digest


@print_time
def remove_background(input_path: str, output_path: str = None, crop_transparent: bool = True) -> str:
    """
    使用 rembg 移除图片背景

    Args:
        input_path: 输入图片路径
        output_path: 输出图片保存路径（可选，如果未提供将自动生成）
        crop_transparent: 是否裁剪透明区域以减小图片大小

    Returns:
        移除背景后的输出图片路径
    """
    try:
        # Auto-generate output path if not provided
        if output_path is None:
            base_path, ext = os.path.splitext(input_path)
            # Force PNG extension to support transparency
            output_path = f"{base_path}_rmbg.png"

        # Remove existing output file if it exists
        if os.path.exists(output_path):
            os.remove(output_path)

        # Load input image and remove background
        input_image = Image.open(input_path)
        output_image = remove(input_image)

        # Ensure the output image has transparency (RGBA mode)
        if output_image.mode != 'RGBA':
            output_image = output_image.convert('RGBA')

        # Crop transparent areas if requested
        if crop_transparent:
            # Get the bounding box of non-transparent pixels
            bbox = output_image.getbbox()
            if bbox:
                output_image = output_image.crop(bbox)

        # Save as PNG to preserve transparency
        output_image.save(output_path, format='PNG')
        return output_path
    except Exception as e:
        raise Exception(f"Failed to remove background from image: {e}")


@print_time
def convert_image_to_svg(input_path: str, output_path: str = None) -> str:
    """
    通过将图片嵌入为 base64 数据将图片转换为 SVG 格式

    Args:
        input_path: 输入图片路径
        output_path: SVG 文件保存路径（可选，如果未提供将自动生成）

    Returns:
        输出 SVG 文件的路径
    """
    try:
        # Auto-generate output path if not provided
        if output_path is None:
            base_path = os.path.splitext(input_path)[0]
            output_path = f"{base_path}.svg"

        # Remove existing output file if it exists
        if os.path.exists(output_path):
            os.remove(output_path)

        # Load and process the image
        with Image.open(input_path) as img:
            # Convert to RGB if necessary (for PNG with transparency)
            if img.mode in ('RGBA', 'LA'):
                # Create white background
                background = Image.new('RGB', img.size, (255, 255, 255))
                if img.mode == 'RGBA':
                    background.paste(img, mask=img.split()[-1])  # Use alpha channel as mask
                else:
                    background.paste(img, mask=img.split()[-1])  # Use alpha channel as mask
                img = background
            elif img.mode != 'RGB':
                img = img.convert('RGB')

            # Convert image to base64
            buffer = BytesIO()
            img.save(buffer, format='PNG')
            img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')

            # Get image dimensions
            width, height = img.size

            # Create SVG content
            svg_content = f'''<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink"
     width="{width}" height="{height}" viewBox="0 0 {width} {height}">
    <image x="0" y="0" width="{width}" height="{height}"
           xlink:href="data:image/png;base64,{img_base64}"/>
</svg>'''

            # Save SVG file
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(svg_content)

            logger.info(f"Image converted to SVG: {output_path}")
            return output_path

    except Exception as e:
        logger.error(f"Failed to convert image to SVG: {str(e)}")
        raise Exception(f"Failed to convert image to SVG: {e}")

def get_image_base64(image_path: str) -> Optional[str]:
    """
    获取图片的base64编码，兼容本地文件和网络图片

    Args:
        image_path: 图片路径，可以是本地文件路径或网络URL

    Returns:
        图片的base64编码字符串，失败时返回None
    """
    try:
        if image_path.startswith(('http://', 'https://')):
            # 处理网络图片
            response = requests.get(image_path, timeout=30)
            response.raise_for_status()
            image_data = response.content
        else:
            # 处理本地文件
            if not os.path.exists(image_path):
                logger.error(f"Local image file not found: {image_path}")
                return None

            with open(image_path, 'rb') as f:
                image_data = f.read()

        # 转换为base64
        base64_str = base64.b64encode(image_data).decode('utf-8')
        return base64_str

    except requests.RequestException as e:
        logger.error(f"Failed to download image from URL {image_path}: {str(e)}")
        return None
    except FileNotFoundError:
        logger.error(f"Local image file not found: {image_path}")
        return None
    except Exception as e:
        logger.error(f"Failed to get base64 for image {image_path}: {str(e)}")
        return None

def get_image_prefix(base64_data: str = '') -> str:
    """获取图像数据的前缀"""
    image_type = detect_image_type(base64_data)
    return f"data:image/{image_type};base64,"


def detect_image_type(base64_data: str) -> str:
    """检测图像类型"""
    try:
        try:
            decoded = base64.b64decode(base64_data[:20])
        except Exception:
            logger.error('detect_image_type error')
            return 'jpeg'

        if len(decoded) < 4:
            return 'jpeg'

        # 读取前4个字节,比对 Magic Number
        header = decoded[:4]
        header_hex = header.hex().upper()

        if header_hex.startswith('FFD8'):
            return 'jpeg'
        elif header_hex.startswith('89504E47'):
            return 'png'
        elif header_hex.startswith('47494638'):
            return 'gif'
        elif decoded.startswith(b'RIFF') and len(decoded) > 8 and decoded[8:12] == b'WEBP':
            return 'webp'
    except Exception:
        logger.error('detect_image_type error')
    return 'jpeg'
